package com.tfswx.ocrengine.util;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * 讯飞识别结果工具类
 *
 * @author Ricky
 * @date 2021/1/11
 */
public class XunfeiOcrUtils {
    private XunfeiOcrUtils() {
    }

    private static final String FIELD_TEXT_LINE = "textline";
    private static final String FIELD_VALUE = "value";

    /**
     * 将讯飞的OCR识别结果装换为阿里的OCR识别结果
     *
     * @param dataSource
     * @return
     */
    public static JSONObject convert2Ali(JSONObject dataSource) {
        JSONObject body = dataSource.getJSONObject("body");
        JSONObject taskResult = body.getJSONObject("task_result");
        JSONArray pages = taskResult.getJSONArray("page");

        JSONObject aliDataSource = new JSONObject();
        aliDataSource.put("prism_version", body.getString("Version"));

        List<JSONObject> wordsInfo = new ArrayList<>();
        List<JSONObject> prismTableInfo = new ArrayList<>();

        StringBuilder contentBuilder = new StringBuilder();
        if (pages != null && !pages.isEmpty()) {
            for (int i = 0; i < pages.size(); i++) {
                JSONObject page = pages.getJSONObject(i);
                if (page == null) {
                    continue;
                }

                JSONObject rect = page.getJSONObject("rect");

                if (rect == null) {
                    continue;
                }

                parasTextRegion(wordsInfo, contentBuilder, page);

                parasTable(wordsInfo, contentBuilder, page);

                parasTableInfos(prismTableInfo, page);
                aliDataSource.put("height", rect.getIntValue("h"));
                aliDataSource.put("width", rect.getIntValue("w"));
                aliDataSource.put("orgHeight", rect.getIntValue("h"));
                aliDataSource.put("orgWidth", rect.getIntValue("w"));
            }
            aliDataSource.put("angle", pages.getJSONObject(0).getInteger("orientation"));
        }
        aliDataSource.put("prism_wordsInfo", wordsInfo);
        aliDataSource.put("prism_tablesInfo", prismTableInfo);

        aliDataSource.put("content", contentBuilder.toString());
        aliDataSource.put("prism_wnum", wordsInfo.size());

        return aliDataSource;
    }

    private static void parasTableInfos(List<JSONObject> prismTableInfo, JSONObject page) {
        JSONObject tableInfo;
        JSONArray tables = page.getJSONArray("Table");
        if (tables == null || tables.isEmpty()) {
            return;
        }
        for (int j = 0; j < tables.size(); j++) {
            tableInfo = new JSONObject();
            JSONObject table = tables.getJSONObject(j);

            tableInfo.put("tableId", table.getInteger("id"));
            tableInfo.put("xCellSize", 1);
            tableInfo.put("yCellSize", 1);

            JSONArray cellInfos = new JSONArray();
            JSONObject cellInfo = null;
            JSONArray cells = JSON.parseArray(table.getString("cell"));
            if (cells != null && !cells.isEmpty()) {
                for (int k = 0; k < cells.size(); k++) {
                    JSONObject cell = cells.getJSONObject(k);
                    cellInfo = getCellInfo(cell);
                    if (cellInfo == null) {
                        continue;
                    }
                    cellInfos.add(cellInfo);
                }
                tableInfo.put("cellInfos", cellInfos);
                prismTableInfo.add(tableInfo);
            }
        }

    }

    private static JSONObject getCellInfo(JSONObject cell) {

        JSONObject cellInfo = new JSONObject();
        cellInfo.put("tableCellId", 0);
        cellInfo.put("xsc", cell.getInteger("col"));
        cellInfo.put("xec", cell.getInteger("col") + cell.getInteger("colspan") - 1);
        cellInfo.put("ysc", cell.getInteger("row"));
        cellInfo.put("yec", cell.getInteger("row") + cell.getInteger("rowspan") - 1);


        JSONObject rect = cell.getJSONObject("rect");
        JSONArray pos = new JSONArray();
        HashMap<String, Integer> hashMap = new HashMap<>(2);
        hashMap.put("x", rect.getInteger("x"));
        hashMap.put("y", rect.getInteger("y"));
        pos.add(hashMap);
        HashMap<String, Integer> hashMap1 = new HashMap<>(2);
        hashMap1.put("x", rect.getInteger("x") + rect.getInteger("w"));
        hashMap1.put("y", rect.getInteger("y"));
        pos.add(hashMap1);
        HashMap<String, Integer> hashMap2 = new HashMap<>(2);
        hashMap2.put("x", rect.getInteger("x"));
        hashMap2.put("y", rect.getInteger("y") + rect.getInteger("h"));
        pos.add(hashMap2);
        HashMap<String, Integer> hashMap3 = new HashMap<>(2);
        hashMap3.put("x", rect.getInteger("x") + rect.getInteger("w"));
        hashMap3.put("y", rect.getInteger("y") + rect.getInteger("h"));
        pos.add(hashMap3);
        cellInfo.put("pos", pos);
        StringBuilder sb = new StringBuilder();
        JSONArray textlines = cell.getJSONArray(FIELD_TEXT_LINE);
        if (textlines == null || textlines.isEmpty()) {
            return null;
        }

        for (int l = 0; l < textlines.size(); l++) {
            JSONArray sents = textlines.getJSONObject(l).getJSONArray("sent");
            if (sents == null || sents.isEmpty()) {
                continue;
            }
            for (int m = 0; m < sents.size(); m++) {
                sb.append(sents.getJSONObject(m).getString(FIELD_VALUE));
            }
        }
        cellInfo.put("word", sb.toString());
        return cellInfo;
    }

    /**
     * 解析识别结果中的表格信息
     *
     * @param wordsInfo
     * @param contentBuilder
     * @param page
     */
    private static void parasTable(List<JSONObject> wordsInfo, StringBuilder contentBuilder, JSONObject page) {
        JSONArray tableArray = page.getJSONArray("Table");
        if (tableArray != null) {
            for (int i1 = 0; i1 < tableArray.size(); i1++) {
                JSONArray cellArray = tableArray.getJSONObject(i1).getJSONArray("cell");
                if (cellArray != null) {
                    for (int i2 = 0; i2 < cellArray.size(); i2++) {
                        JSONArray textlines = cellArray.getJSONObject(i2).getJSONArray(FIELD_TEXT_LINE);
                        if (textlines != null) {
                            parasTextLine(wordsInfo, contentBuilder, textlines);
                        }
                    }
                }
            }
        }
    }

    /**
     * 解析识别结果中的文本信息
     *
     * @param wordsInfo
     * @param contentBuilder
     * @param page
     */
    private static void parasTextRegion(List<JSONObject> wordsInfo, StringBuilder contentBuilder, JSONObject page) {
        // 处理文本识别结果
        JSONArray textRegion = page.getJSONArray("TextRegion");
        if (textRegion != null) {
            for (int j = 0; j < textRegion.size(); j++) {
                JSONArray textline = textRegion.getJSONObject(j).getJSONArray(FIELD_TEXT_LINE);
                if (textline != null) {
                    parasTextLine(wordsInfo, contentBuilder, textline);
                }
            }
        }
    }

    /**
     * 解析讯飞OCR识别结果中的TextLine信息
     *
     * @param wordsInfo
     * @param contentBuilder
     * @param textline
     */
    private static void parasTextLine(List<JSONObject> wordsInfo, StringBuilder contentBuilder, JSONArray textline) {
        for (int k = 0; k < textline.size(); k++) {
            JSONObject textLineJsonObject = textline.getJSONObject(k);
            JSONArray sent = textLineJsonObject.getJSONArray("sent");
            if (sent == null) {
                continue;
            }
            for (int m = 0; m < sent.size(); m++) {
                JSONObject sentItem = sent.getJSONObject(m);
                if (sentItem == null) {
                    continue;
                }
                JSONObject wordInfo = getWordInfo(textLineJsonObject, sentItem, contentBuilder, m);
                wordsInfo.add(wordInfo);

            }

        }
    }

    private static JSONObject getWordInfo(JSONObject textLineJsonObject, JSONObject sentItem, StringBuilder contentBuilder, int rowId) {

        JSONObject wordInfo = new JSONObject();

        wordInfo.put("word", sentItem.getString(FIELD_VALUE));
        wordInfo.put("prob", 99);

        JSONArray charArr = sentItem.getJSONArray("char");
        if (charArr != null) {
            List<JSONObject> charInfo = new ArrayList<>();
            for (int l = 0; l < charArr.size(); l++) {
                JSONObject char1 = charArr.getJSONObject(l);
                JSONObject rect = JSON.parseObject(char1.get("rect").toString());

                JSONObject tlPoint = JSON.parseObject(rect.get("tl_point").toString());
                JSONObject trPoint = JSON.parseObject(rect.get("tr_point").toString());
                JSONObject blPoint = JSON.parseObject(rect.get("bl_point").toString());

                JSONObject word = new JSONObject();
                word.put("word", char1.get(FIELD_VALUE));
                word.put("prob", 99);
                word.put("x", tlPoint.get("x"));
                word.put("y", tlPoint.get("y"));
                word.put("w", Integer.parseInt(trPoint.get("x").toString()) - Integer.parseInt(tlPoint.get("x").toString()));
                word.put("h", Integer.parseInt(blPoint.get("y").toString()) - Integer.parseInt(tlPoint.get("y").toString()));
                charInfo.add(word);
            }
            wordInfo.put("charInfo", charInfo);
        } else {
            wordInfo.put("charInfo", new ArrayList<>(0));
        }

        JSONArray posArray = new JSONArray();
        posArray.add(textLineJsonObject.getJSONObject("tl_point"));
        posArray.add(textLineJsonObject.getJSONObject("tr_point"));
        posArray.add(textLineJsonObject.getJSONObject("br_point"));
        posArray.add(textLineJsonObject.getJSONObject("bl_point"));

        wordInfo.put("pos", posArray);

        contentBuilder.append(sentItem.getString(FIELD_VALUE)).append("\n");
        wordInfo.put("rowId", rowId);
        wordInfo.put("paragraphId", rowId);
        return wordInfo;
    }
}
